import numpy as np
import treelib

from generator.trajectory_generator import Generator
from config.parameter_carrier import ParameterCarrier
from pritree.tree_model import TreeModel


class StateGeneration:

    def __init__(self, cc: ParameterCarrier):
        self.cc = cc

    def generate_tra(self, mar_mo):
        cc1 = self.cc
        generator1 = Generator(self.cc)
        generator1.load_generator(mar_mo)
        number = cc1.trajectory_number_to_generate
        usable_tr_list = generator1.generate_many(number, neighbor_check=False)  # 生成网格的轨迹
        print('state trajectories got')
        return usable_tr_list

    ########树生成过程
    def generate_tra_by_tree(self, mar_mo, tm: TreeModel):
        cc1 = self.cc
        generator1 = Generator(self.cc)
        generator1.load_generator(mar_mo)
        number = cc1.trajectory_number_to_generate
        usable_tr_list = generator1.generate_many_by_tree(number, tm)  # 生成网格的轨迹
        print('state trajectories got')

        return usable_tr_list

    # this function transfers a list of usable state trajectory to real state trajectory
